
import torch
import copy
from torch import nn, optim
from torch.utils import data 


import logging
from trainers.fedbase.FedClientBase import FedClientBase

# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format ='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)


class FedProxClient(FedClientBase):
    """训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。\n
    这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。
    1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
    2. 训练的模型、参数在客户端配置。
    3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。
    
    接收参数，完成训练。主要是一次训练的参数。
    每一个算法都应该提供梯度的返回值。即在客户端计算出梯度。其实在服务器计算也无可厚非。
    """    
 
    def __init__(self,datapair,model):
        super(FedProxClient, self).__init__(datapair,model)
        logger.info("FedProxClient init----:")


    def set_parameters(self,client_params):
        super().set_parameters(client_params)
        # prox本地更新的权重
        if("prox_mu" in client_params):
            self.prox_mu = client_params["prox_mu"]
        else:
            self.prox_mu = 0.0
        
        logger.info("FedProxClient set_parameters:prox_mu={}".format(self.prox_mu))
 

    def train(self):
        """ 用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy。\n
        添加proximate update。在每次更新完成后。
        1. 直接更改优化器比较麻烦，尝试对loss_func进行更改。看一下是否能够达到要求
        2. 尝试直接修改权重参数
        3. 尝试直接继承修改后的优化器

        Returns:
            [tuple(loss,accuracy)]: 损失和训练准确度
        """        
        # 得到最开始parmlist，作为限制权重的基准
        old_parm_list = []
        for parm in self.model.parameters():
            old_parm_list.append(torch.clone(parm).detach())   
        

        # param_state['old_init'] = torch.clone(p.data).detach()
        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0
        for epoch in range(self.epochs):
            for batch_x,batch_y in self.train_dl:
                # 正向传播计算准确度
                outputs = self.model(batch_x)
                _,predicted = torch.max(outputs,1)
                correct_sum += (predicted == batch_y).sum().item()
                total_size += len(batch_y)
                loss = self.loss_func(outputs,batch_y.long())

                # 利用梯度下降过程，完成prox正则化项的计算。
                # 经过验证发现通过+=构造的连续的计算图可以正常进行梯度下降
                for i,parm in enumerate(self.model.parameters()):
                    loss += (1.0/2.0*self.prox_mu * (parm-old_parm_list[i])**2).sum()
                
                loss_sum +=loss.item()
                loss_times+=1
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()
                # 梯度下降
                self.optimizer.step()
        # print(self.optimizer.state_dict())
        return loss_sum/loss_times,100*correct_sum/total_size

    # 用来测试训练的结果返回准确率
    def test(self):
        self.model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.train_dl:
                outputs = self.model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy
